import sqlalchemy
from k_means_constrained import KMeansConstrained
import numpy as np

def get_dataset(class_name: str, engine):
    conn = engine.connect()
    try:
        result = conn.execute(sqlalchemy.text("""
                            SELECT stu_detail.uid, stu_detail.stu_num, user_info.username
                            FROM stu_detail, user_info 
                            WHERE user_info.uid=stu_detail.uid AND stu_detail.class_name=:class_name
                        """), {
            "class_name": class_name
        })
        stu_list = result.fetchall()
        grade = {stu[0]: [] for stu in stu_list}
        info_dict = {stu[0]: {'uid': stu[0], 'stu_num': stu[1], 'username': stu[2]} for stu in stu_list}

        # 查询班级中至少有一个学生有成绩的课程
        result = conn.execute(sqlalchemy.text("""
                            SELECT DISTINCT lesson_num
                            FROM stu_grade
                            WHERE uid IN (
                                SELECT uid FROM stu_detail WHERE class_name = :class_name
                            )
                        """), {
            "class_name": class_name
        })
        lesson_list = result.fetchall()

        for stu in stu_list:
            for lesson in lesson_list:
                # 查询学生在该课程的成绩
                score_result = conn.execute(sqlalchemy.text("""
                                           SELECT grade
                                           FROM stu_grade
                                           WHERE uid = :uid AND lesson_num = :lesson_num
                                       """), {
                    "uid": stu[0],
                    "lesson_num": lesson[0]
                })
                score = score_result.fetchone()
                if score:
                    grade[stu[0]].append(float(score[0]))
                else:
                    grade[stu[0]].append(0.0)

        conn.commit()
    except Exception as e:
        print(e)
        raise e
    finally:
        conn.close()

    return grade, info_dict


def group_for_class(class_name: str, num_of_groups: int, engine):
    dataset, stu_info_dict = get_dataset(class_name, engine)
    print(dataset, stu_info_dict)
    # {'18988881001': [78.0, 61.0, 45.0, 65.0, 60.0, 81.0], '18988881002': [98.0, 98.0, 93.0, 80.0, 95.0, 98.0], '18988881003': [86.0, 60.0, 94.0, 53.0, 54.0, 86.0], '18988881004': [48.0, 94.0, 74.0, 99.0, 90.0, 50.0], '18988881006': [64.0, 44.0, 86.0, 68.0, 85.0, 66.0]}
    # 将数据转换为numpy数组
    data = np.array(list(dataset.values()))
    if num_of_groups > len(data):
        return None

    # 使用KMeansConstrained进行聚类
    clf = KMeansConstrained(
        n_clusters=num_of_groups,
        size_min=max(1, len(data) // num_of_groups),
        size_max=len(data) // num_of_groups+1,
        random_state=0
    )

    # 拟合数据并预测
    labels = clf.fit_predict(data)

    # 创建分组字典
    groups = {i: [] for i in range(num_of_groups)}

    # 将学生信息分配到对应的组
    for idx, label in enumerate(labels):
        phone = list(dataset.keys())[idx]
        groups[label].append(stu_info_dict[phone])

    # # 打印分组结果
    # for group_id, members in groups.items():
    #     print(f"Group {group_id + 1}:")
    #     for member in members:
    #         print(member)
    print(groups)
    # {0: [{'stu_num': '21019001', 'username': '张伟'}, {'stu_num': '21018003', 'username': '王强'}, {'stu_num': '21019006', 'username': '徐洋涛'}], 1: [{'stu_num': '21019002', 'username': '李娜'}, {'stu_num': '21019004', 'username': '赵敏'}]}
    return groups

